#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
告警管理系统

提供多渠道告警通知功能，包括邮件、微信、短信等。
"""

import asyncio
import logging
import json
import smtplib
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Callable
from dataclasses import dataclass, asdict
from enum import Enum
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
import requests

class AlertChannel(Enum):
    """告警渠道"""
    EMAIL = "email"
    WECHAT = "wechat"
    SMS = "sms"
    WEBHOOK = "webhook"
    CONSOLE = "console"

class AlertStatus(Enum):
    """告警状态"""
    PENDING = "pending"
    SENT = "sent"
    FAILED = "failed"
    RESOLVED = "resolved"

@dataclass
class AlertConfig:
    """告警配置"""
    channel: AlertChannel
    enabled: bool = True
    config: Dict = None
    
    def __post_init__(self):
        if self.config is None:
            self.config = {}

@dataclass
class Alert:
    """告警信息"""
    id: str
    rule_id: str
    rule_name: str
    type: str
    severity: str
    title: str
    content: str
    target_id: str
    timestamp: str
    status: AlertStatus = AlertStatus.PENDING
    channels: List[AlertChannel] = None
    retry_count: int = 0
    last_retry: str = ""
    resolved_time: str = ""
    
    def __post_init__(self):
        if self.channels is None:
            self.channels = [AlertChannel.CONSOLE]

class AlertManager:
    """告警管理器"""
    
    def __init__(self, config: Dict = None):
        self.logger = logging.getLogger(__name__)
        
        # 默认配置
        self.config = config or self._get_default_config()
        
        # 告警渠道配置
        self.channel_configs: Dict[AlertChannel, AlertConfig] = {}
        self._init_channel_configs()
        
        # 告警处理器
        self.channel_handlers: Dict[AlertChannel, Callable] = {
            AlertChannel.EMAIL: self._send_email_alert,
            AlertChannel.WECHAT: self._send_wechat_alert,
            AlertChannel.SMS: self._send_sms_alert,
            AlertChannel.WEBHOOK: self._send_webhook_alert,
            AlertChannel.CONSOLE: self._send_console_alert
        }
        
        # 告警历史和统计
        self.alert_history: List[Alert] = []
        self.stats = {
            "total_alerts": 0,
            "sent_alerts": 0,
            "failed_alerts": 0,
            "channel_stats": {},
            "severity_stats": {}
        }
        
        self.logger.info("告警管理器初始化完成")
    
    def _get_default_config(self) -> Dict:
        """获取默认配置"""
        return {
            "max_retry_count": 3,
            "retry_interval": 300,  # 秒
            "alert_cooldown": 600,  # 秒
            "batch_size": 10,
            "enable_console_fallback": True,
            "history_size": 1000
        }
    
    def _init_channel_configs(self):
        """初始化渠道配置"""
        # 邮件配置
        self.channel_configs[AlertChannel.EMAIL] = AlertConfig(
            channel=AlertChannel.EMAIL,
            enabled=False,  # 默认禁用，需要用户配置
            config={
                "smtp_server": "",
                "smtp_port": 587,
                "username": "",
                "password": "",
                "from_email": "",
                "to_emails": [],
                "use_tls": True
            }
        )
        
        # 微信配置
        self.channel_configs[AlertChannel.WECHAT] = AlertConfig(
            channel=AlertChannel.WECHAT,
            enabled=False,
            config={
                "webhook_url": "",
                "secret": ""
            }
        )
        
        # 短信配置
        self.channel_configs[AlertChannel.SMS] = AlertConfig(
            channel=AlertChannel.SMS,
            enabled=False,
            config={
                "api_url": "",
                "api_key": "",
                "phone_numbers": []
            }
        )
        
        # Webhook配置
        self.channel_configs[AlertChannel.WEBHOOK] = AlertConfig(
            channel=AlertChannel.WEBHOOK,
            enabled=False,
            config={
                "url": "",
                "method": "POST",
                "headers": {},
                "auth": {}
            }
        )
        
        # 控制台配置（总是启用）
        self.channel_configs[AlertChannel.CONSOLE] = AlertConfig(
            channel=AlertChannel.CONSOLE,
            enabled=True,
            config={}
        )
    
    async def send_alert(self, alert_data: Dict) -> bool:
        """发送告警"""
        # 创建告警对象
        alert = Alert(
            id=f"alert_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{len(self.alert_history)}",
            rule_id=alert_data.get('rule_id', ''),
            rule_name=alert_data.get('rule_name', ''),
            type=alert_data.get('type', ''),
            severity=alert_data.get('severity', 'info'),
            title=alert_data.get('title', ''),
            content=alert_data.get('content', ''),
            target_id=alert_data.get('target_id', ''),
            timestamp=alert_data.get('timestamp', datetime.now().isoformat()),
            channels=self._determine_channels(alert_data.get('severity', 'info'))
        )
        
        # 检查告警冷却
        if self._is_in_cooldown(alert):
            self.logger.info(f"告警在冷却期内，跳过: {alert.title}")
            return True
        
        success = await self._process_alert(alert)
        
        # 保存到历史记录
        self.alert_history.append(alert)
        if len(self.alert_history) > self.config["history_size"]:
            self.alert_history = self.alert_history[-self.config["history_size"]:]
        
        # 更新统计信息
        self._update_stats(alert)
        
        return success
    
    def _determine_channels(self, severity: str) -> List[AlertChannel]:
        """根据严重程度确定告警渠道"""
        if severity == "critical":
            # 严重告警使用所有可用渠道
            return [channel for channel, config in self.channel_configs.items() 
                   if config.enabled]
        elif severity == "high":
            # 高级告警使用邮件和微信
            return [AlertChannel.EMAIL, AlertChannel.WECHAT, AlertChannel.CONSOLE]
        elif severity in ["medium", "warning"]:
            # 中级告警使用邮件
            return [AlertChannel.EMAIL, AlertChannel.CONSOLE]
        else:
            # 低级别告警只使用控制台
            return [AlertChannel.CONSOLE]
    
    def _is_in_cooldown(self, alert: Alert) -> bool:
        """检查告警是否在冷却期"""
        cooldown_seconds = self.config["alert_cooldown"]
        cutoff_time = datetime.now() - timedelta(seconds=cooldown_seconds)
        
        # 检查相同类型和目标的告警
        for historical_alert in reversed(self.alert_history):
            if (historical_alert.type == alert.type and 
                historical_alert.target_id == alert.target_id and
                historical_alert.status == AlertStatus.SENT):
                
                alert_time = datetime.fromisoformat(historical_alert.timestamp)
                if alert_time > cutoff_time:
                    return True
                break
        
        return False
    
    async def _process_alert(self, alert: Alert) -> bool:
        """处理单个告警"""
        success_channels = []
        failed_channels = []
        
        # 尝试通过各个渠道发送
        for channel in alert.channels:
            if channel not in self.channel_configs:
                continue
            
            config = self.channel_configs[channel]
            if not config.enabled:
                continue
            
            try:
                handler = self.channel_handlers.get(channel)
                if handler:
                    success = await handler(alert, config.config)
                    if success:
                        success_channels.append(channel)
                    else:
                        failed_channels.append(channel)
                else:
                    self.logger.warning(f"未找到渠道处理器: {channel}")
                    failed_channels.append(channel)
                    
            except Exception as e:
                self.logger.error(f"发送告警失败 [{channel.value}]: {e}")
                failed_channels.append(channel)
        
        # 更新告警状态
        if success_channels:
            alert.status = AlertStatus.SENT
            self.logger.info(
                f"告警发送成功: {alert.title}, "
                f"成功渠道: {[ch.value for ch in success_channels]}"
            )
            
            # 如果有失败的渠道，记录重试信息
            if failed_channels and alert.retry_count < self.config["max_retry_count"]:
                alert.last_retry = datetime.now().isoformat()
                alert.retry_count += 1
                # 可以实现重试逻辑
            
            return True
        else:
            alert.status = AlertStatus.FAILED
            self.logger.error(f"告警发送失败: {alert.title}")
            
            # 如果启用控制台回退且控制台不在原始渠道中
            if (self.config["enable_console_fallback"] and 
                AlertChannel.CONSOLE not in alert.channels):
                await self._send_console_alert(alert, {})
            
            return False
    
    async def _send_email_alert(self, alert: Alert, config: Dict) -> bool:
        """发送邮件告警"""
        try:
            if not all([config.get('smtp_server'), config.get('username'), 
                       config.get('password'), config.get('from_email')]):
                self.logger.warning("邮件配置不完整")
                return False
            
            # 创建邮件内容
            msg = MIMEMultipart()
            msg['From'] = config['from_email']
            msg['To'] = ", ".join(config.get('to_emails', []))
            msg['Subject'] = f"[{alert.severity.upper()}] {alert.title}"
            
            # 邮件正文
            body = self._format_email_body(alert)
            msg.attach(MIMEText(body, 'html', 'utf-8'))
            
            # 发送邮件
            server = smtplib.SMTP(config['smtp_server'], config.get('smtp_port', 587))
            
            if config.get('use_tls', True):
                server.starttls()
            
            server.login(config['username'], config['password'])
            server.send_message(msg)
            server.quit()
            
            return True
            
        except Exception as e:
            self.logger.error(f"发送邮件告警失败: {e}")
            return False
    
    def _format_email_body(self, alert: Alert) -> str:
        """格式化邮件正文"""
        return f"""
        <html>
        <body>
            <h2>舆情监控告警</h2>
            <table border="1" style="border-collapse: collapse; width: 100%;">
                <tr><td><strong>告警标题</strong></td><td>{alert.title}</td></tr>
                <tr><td><strong>严重程度</strong></td><td><span style="color: {'red' if alert.severity == 'critical' else 'orange' if alert.severity == 'high' else 'blue'};">{alert.severity.upper()}</span></td></tr>
                <tr><td><strong>告警类型</strong></td><td>{alert.type}</td></tr>
                <tr><td><strong>触发规则</strong></td><td>{alert.rule_name}</td></tr>
                <tr><td><strong>告警时间</strong></td><td>{alert.timestamp}</td></tr>
                <tr><td><strong>监控目标ID</strong></td><td>{alert.target_id}</td></tr>
            </table>
            <h3>详细信息：</h3>
            <p>{alert.content.replace(chr(10), '<br>')}</p>
            <hr>
            <small>此邮件由舆情监控系统自动发送</small>
        </body>
        </html>
        """
    
    async def _send_wechat_alert(self, alert: Alert, config: Dict) -> bool:
        """发送微信告警"""
        try:
            webhook_url = config.get('webhook_url')
            if not webhook_url:
                self.logger.warning("微信webhook配置不完整")
                return False
            
            # 构造消息内容
            severity_emoji = {
                'critical': '🚨',
                'high': '⚠️',
                'medium': '⚡',
                'low': 'ℹ️',
                'info': 'ℹ️'
            }
            
            emoji = severity_emoji.get(alert.severity, 'ℹ️')
            
            message = {
                "msgtype": "markdown",
                "markdown": {
                    "content": f"""# {emoji} 舆情监控告警
**告警标题:** {alert.title}
**严重程度:** <font color="{self._get_severity_color(alert.severity)}">{alert.severity.upper()}</font>
**告警类型:** {alert.type}
**触发规则:** {alert.rule_name}
**告警时间:** {alert.timestamp}
**监控目标:** {alert.target_id}

**详细信息:**
{alert.content}
"""
                }
            }
            
            # 发送请求
            response = requests.post(webhook_url, json=message, timeout=10)
            return response.status_code == 200
            
        except Exception as e:
            self.logger.error(f"发送微信告警失败: {e}")
            return False
    
    def _get_severity_color(self, severity: str) -> str:
        """获取严重程度对应的颜色"""
        colors = {
            'critical': 'red',
            'high': 'orange',
            'medium': 'blue',
            'low': 'green',
            'info': 'gray'
        }
        return colors.get(severity, 'gray')
    
    async def _send_sms_alert(self, alert: Alert, config: Dict) -> bool:
        """发送短信告警"""
        try:
            api_url = config.get('api_url')
            api_key = config.get('api_key')
            phone_numbers = config.get('phone_numbers', [])
            
            if not all([api_url, api_key, phone_numbers]):
                self.logger.warning("短信配置不完整")
                return False
            
            # 构造短信内容（简化版）
            content = f"[{alert.severity.upper()}] {alert.title}\\n{alert.content[:100]}..."
            
            # 发送短信（示例实现，需要根据具体短信服务商调整）
            for phone in phone_numbers:
                data = {
                    "api_key": api_key,
                    "mobile": phone,
                    "content": content
                }
                
                response = requests.post(api_url, json=data, timeout=10)
                if response.status_code != 200:
                    return False
            
            return True
            
        except Exception as e:
            self.logger.error(f"发送短信告警失败: {e}")
            return False
    
    async def _send_webhook_alert(self, alert: Alert, config: Dict) -> bool:
        """发送Webhook告警"""
        try:
            url = config.get('url')
            if not url:
                self.logger.warning("Webhook配置不完整")
                return False
            
            method = config.get('method', 'POST')
            headers = config.get('headers', {})
            auth = config.get('auth', {})
            
            # 构造请求数据
            data = {
                "alert_id": alert.id,
                "rule_id": alert.rule_id,
                "rule_name": alert.rule_name,
                "type": alert.type,
                "severity": alert.severity,
                "title": alert.title,
                "content": alert.content,
                "target_id": alert.target_id,
                "timestamp": alert.timestamp
            }
            
            # 发送请求
            if method.upper() == 'POST':
                response = requests.post(url, json=data, headers=headers, 
                                       auth=auth if auth else None, timeout=10)
            else:
                response = requests.get(url, params=data, headers=headers,
                                      auth=auth if auth else None, timeout=10)
            
            return response.status_code in [200, 201, 202]
            
        except Exception as e:
            self.logger.error(f"发送Webhook告警失败: {e}")
            return False
    
    async def _send_console_alert(self, alert: Alert, config: Dict) -> bool:
        """发送控制台告警"""
        try:
            severity_colors = {
                'critical': '\\033[91m',  # 红色
                'high': '\\033[93m',      # 黄色
                'medium': '\\033[94m',    # 蓝色
                'low': '\\033[92m',       # 绿色
                'info': '\\033[0m'        # 默认
            }
            
            color = severity_colors.get(alert.severity, '\\033[0m')
            reset = '\\033[0m'
            
            message = f"""
{color}=============== 舆情监控告警 ==============={reset}
告警ID: {alert.id}
告警标题: {color}{alert.title}{reset}
严重程度: {color}{alert.severity.upper()}{reset}
告警类型: {alert.type}
触发规则: {alert.rule_name}
告警时间: {alert.timestamp}
监控目标: {alert.target_id}

详细信息:
{alert.content}
{color}=========================================={reset}
"""
            
            print(message)
            self.logger.warning(f"[ALERT] {alert.title} - {alert.severity}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"发送控制台告警失败: {e}")
            return False
    
    def configure_channel(self, channel: AlertChannel, config: Dict, 
                         enabled: bool = True) -> bool:
        """配置告警渠道"""
        try:
            if channel not in self.channel_configs:
                self.channel_configs[channel] = AlertConfig(channel=channel)
            
            self.channel_configs[channel].config.update(config)
            self.channel_configs[channel].enabled = enabled
            
            self.logger.info(f"已配置告警渠道: {channel.value}")
            return True
        except Exception as e:
            self.logger.error(f"配置告警渠道失败: {e}")
            return False
    
    def test_channel(self, channel: AlertChannel) -> bool:
        """测试告警渠道"""
        test_alert = Alert(
            id="test_alert",
            rule_id="test_rule",
            rule_name="测试规则",
            type="test",
            severity="info",
            title="告警渠道测试",
            content="这是一条测试告警消息，用于验证告警渠道配置是否正确。",
            target_id="test_target",
            timestamp=datetime.now().isoformat(),
            channels=[channel]
        )
        
        try:
            # 同步调用异步方法进行测试
            loop = asyncio.get_event_loop()
            success = loop.run_until_complete(self._process_alert(test_alert))
            
            if success:
                self.logger.info(f"告警渠道测试成功: {channel.value}")
            else:
                self.logger.error(f"告警渠道测试失败: {channel.value}")
            
            return success
        except Exception as e:
            self.logger.error(f"测试告警渠道时出错: {e}")
            return False
    
    def get_channel_configs(self) -> Dict:
        """获取所有渠道配置"""
        return {
            channel.value: {
                "enabled": config.enabled,
                "config": config.config
            }
            for channel, config in self.channel_configs.items()
        }
    
    def get_alert_history(self, limit: int = 50) -> List[Dict]:
        """获取告警历史"""
        recent_alerts = self.alert_history[-limit:] if limit > 0 else self.alert_history
        return [asdict(alert) for alert in reversed(recent_alerts)]
    
    def get_statistics(self) -> Dict:
        """获取告警统计信息"""
        # 更新统计信息
        total = len(self.alert_history)
        sent = sum(1 for alert in self.alert_history if alert.status == AlertStatus.SENT)
        failed = sum(1 for alert in self.alert_history if alert.status == AlertStatus.FAILED)
        
        # 按渠道统计
        channel_stats = {}
        for alert in self.alert_history:
            for channel in alert.channels:
                channel_name = channel.value
                if channel_name not in channel_stats:
                    channel_stats[channel_name] = {"sent": 0, "failed": 0}
                
                if alert.status == AlertStatus.SENT:
                    channel_stats[channel_name]["sent"] += 1
                elif alert.status == AlertStatus.FAILED:
                    channel_stats[channel_name]["failed"] += 1
        
        # 按严重程度统计
        severity_stats = {}
        for alert in self.alert_history:
            severity = alert.severity
            severity_stats[severity] = severity_stats.get(severity, 0) + 1
        
        return {
            "total_alerts": total,
            "sent_alerts": sent,
            "failed_alerts": failed,
            "success_rate": sent / max(total, 1),
            "channel_stats": channel_stats,
            "severity_stats": severity_stats,
            "recent_24h": self._get_recent_count(24),
            "recent_1h": self._get_recent_count(1)
        }
    
    def _get_recent_count(self, hours: int) -> int:
        """获取最近N小时的告警数量"""
        cutoff_time = datetime.now() - timedelta(hours=hours)
        
        return sum(1 for alert in self.alert_history 
                  if datetime.fromisoformat(alert.timestamp) > cutoff_time)
    
    def _update_stats(self, alert: Alert):
        """更新统计信息"""
        self.stats["total_alerts"] += 1
        
        if alert.status == AlertStatus.SENT:
            self.stats["sent_alerts"] += 1
        elif alert.status == AlertStatus.FAILED:
            self.stats["failed_alerts"] += 1
        
        # 更新渠道统计
        for channel in alert.channels:
            channel_name = channel.value
            if channel_name not in self.stats["channel_stats"]:
                self.stats["channel_stats"][channel_name] = 0
            self.stats["channel_stats"][channel_name] += 1
        
        # 更新严重程度统计
        severity = alert.severity
        if severity not in self.stats["severity_stats"]:
            self.stats["severity_stats"][severity] = 0
        self.stats["severity_stats"][severity] += 1
    
    async def resolve_alert(self, alert_id: str) -> bool:
        """解决告警"""
        for alert in self.alert_history:
            if alert.id == alert_id:
                alert.status = AlertStatus.RESOLVED
                alert.resolved_time = datetime.now().isoformat()
                self.logger.info(f"告警已解决: {alert.title} ({alert_id})")
                return True
        return False
    
    def cleanup_old_alerts(self, days: int = 7):
        """清理旧告警记录"""
        cutoff_time = datetime.now() - timedelta(days=days)
        
        original_count = len(self.alert_history)
        self.alert_history = [
            alert for alert in self.alert_history
            if datetime.fromisoformat(alert.timestamp) > cutoff_time
        ]
        
        removed_count = original_count - len(self.alert_history)
        if removed_count > 0:
            self.logger.info(f"已清理 {removed_count} 条旧告警记录")